from PyQt5.QtWidgets import QWidget,QMessageBox,QFileDialog
from PyQt5.QtGui import QIntValidator
from PyQt5.QtWebEngineWidgets import QWebEngineView
from PyQt5.QtWebChannel import QWebChannel
from PyQt5.QtCore import QUrl,QThread,QObject,pyqtSignal,pyqtProperty
from ui.uiTrainNet import Ui_trainNet
import os.path as op
from nn import train
import json
from multiprocessing import Process,Queue
from tensorflow.keras.callbacks import Callback
from util import *

class TrainWorker(Callback):
    def __init__(self,trainParams,queue):
        Callback.__init__(self)
        self.trainParams=trainParams
        self.init_loss = None
        self.queue=queue

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        loss = logs.get("loss")
        accuracy = logs.get("accuracy")
        if self.init_loss is None:
            self.init_loss = loss
        self.queue.put([epoch,loss,accuracy])

    def work(self):
        train.train(*self.trainParams,self)
        self.queue.put([-1])

class UpdateWorkder(QObject):
    onEpochEnd=pyqtSignal(int,float,float)
    finished=pyqtSignal()
    def __init__(self,queue):
        super(UpdateWorkder,self).__init__()
        self.queue=queue

    def work(self):
        while True:
            data=self.queue.get()
            if data[0]==-1:
                break
            self.data=data
            self.onEpochEnd.emit(data[0],data[1],data[2])
        self.finished.emit()

class WebPageMiddle(QObject):
    textChanged=pyqtSignal(str)
    def __init__(self):
        super().__init__()
        self._text=""

    def getText(self):
        return self._text

    def setText(self,text):
        self._text=text
        self.textChanged.emit(text)
    text=pyqtProperty(str,fget=getText,fset=setText)

class TrainNetWidget(QWidget,Ui_trainNet):
    logTransmit=pyqtSignal(str)
    def __init__(self,parent=None):
        super(TrainNetWidget,self).__init__(parent)
        self.setupUi(self)
        self.initUi()
        self.middle = WebPageMiddle()
        self.channel = QWebChannel()
        self.channel.registerObject("middle", self.middle)
        self.connectSignals()

    def initUi(self):
        self.le_iterTimes.setValidator(QIntValidator())
        self.le_batchSize.setValidator(QIntValidator())
        self.browser = QWebEngineView()
        self.browser.load(QUrl("file:///chart/train.html"))
        self.mainLayout.addWidget(self.browser)

    def emitLog(self,s):
        self.logTransmit.emit(s)

    def connectSignals(self):
        self.pb_samplePath.clicked.connect(self.chooseSamplePath)
        self.pb_modelPath.clicked.connect(self.chooseModelSavePath)
        self.pb_train.clicked.connect(self.beginTrain)

    def chooseSamplePath(self):
        path=QFileDialog.getExistingDirectory(self,"选择样本数据")
        if path!="":
            if not op.exists(op.join(path,'data.npd')):
                QMessageBox.warning(self,"提示","此目录下找不到数据文件'data.npd'",QMessageBox.Ok,QMessageBox.Ok)
            elif not op.exists(op.join(path,'label.npd')):
                QMessageBox.warning(self,"提示","此目录下找不到数据文件'label.npd'",QMessageBox.Ok,QMessageBox.Ok)
            else:
                self.le_samplePath.setText(path)


    def chooseModelSavePath(self):
        path,_=QFileDialog.getSaveFileName(self,"选择模型保存路径")
        if path!="":
            self.le_modelPath.setText(path)

    def beginTrain(self):
        # samplePath='C:/WorkSpace/zhushuizhan/sim/test.npd'
        # modelPath='C:/WorkSpace/zhushuizhan/sim/test'
        samplePath=self.le_samplePath.text()
        modelPath=self.le_modelPath.text()
        iterTimes=self.le_iterTimes.text()
        batchSize=self.le_batchSize.text()
        if samplePath=='' or modelPath=='' or iterTimes=='' or batchSize=='':
            QMessageBox.critical(self,"错误","参数缺失",QMessageBox.Ok,QMessageBox.Ok)
            return
        iterTimes=int(iterTimes)
        batchSize=int(batchSize)
        dataPath=op.join(samplePath,'data.npd')
        labelPath=op.join(samplePath,'label.npd')

        self.training=True
        self.pb_train.setEnabled(False)

        self.browser.page().runJavaScript("reset()")
        self.queue=Queue()
        trainWorker=TrainWorker((dataPath,labelPath,modelPath+'.tf',iterTimes,batchSize),self.queue)
        p=Process(target=trainWorker.work)
        p.start()
        self.updateWorker=UpdateWorkder(self.queue)
        self.updateThread=QThread()
        self.updateWorker.onEpochEnd.connect(self.updateChart)
        self.updateWorker.finished.connect(self.trainFinished)
        self.updateWorker.moveToThread(self.updateThread)
        self.updateThread.started.connect(self.updateWorker.work)
        self.updateThread.start()

    def trainFinished(self):
        self.training=False
        self.pb_train.setEnabled(True)
        self.updateThread.terminate()
        s="训练完成，最终loss值:{}, 精度:{}".format(self.updateWorker.data[1],self.updateWorker.data[2])
        self.emitLog(okLog(s))

    def updateChart(self,epoch,loss,accuracy):
        dic={"轮数":epoch,"loss值":loss,"精度":accuracy}
        text=json.dumps(dic)
        self.middle.text=text

    def resizeEvent(self, event):
        self.browser.page().runJavaScript("resize()")
        super().resizeEvent(event)

    def resizeLater(self,milliseconds):
        self.browser.page().runJavaScript("resizeLater(%d)" % milliseconds)


if __name__=="__main__":
    import sys
    from PyQt5.QtWidgets import QApplication
    app = QApplication(sys.argv)
    exe = TrainNetWidget()
    exe.show()
    exe.browser.page().setWebChannel(exe.channel)
    sys.exit(app.exec_())